import numpy as np
from math import asinh, atan, sqrt, pi
import torch#GPU
import torch.fft
import time
from matplotlib import pyplot as plt
from multiprocessing import Pool,Process
#==========================================================================
class sim:
    def __init__(self,
        name,namedemag,n,dx,m0,gamma,ms   ,A     ,alpha,h_zee,T,T0,
        error=1e-5,mindt=1e-15,maxdt=1e-11,fixdt=0,mu0   = 4e-7 * pi,inscribed=False):
        self.name=name
        self.namedemag=namedemag
        self.n=n
        self.dx=dx
        self.m0=m0
        self.gamma=gamma
        self.ms=ms
        self.A=A
        self.alpha=alpha
        self.h_zee=h_zee
        self.T=T
        self.T0=T0
        self.error=error
        self.mindt=mindt
        self.maxdt=maxdt
        self.fixdt=fixdt
        self.mu0=mu0
        self.inscribed=inscribed
        #=======================
        self.m = torch.zeros(n + [3,],dtype=torch.float32).cuda()
        self.m_pad     =torch.zeros(([2*i-1 for i in n] + [3,]),dtype=torch.float32).cuda()
        self.f_n_demag=torch.zeros([2*i-1 for i in n] + [6,],dtype=torch.complex64).cuda()
        self.m_form=torch.zeros((n + [3,]),dtype=torch.int).cuda()
        self.surfacex=torch.zeros((n + [3,]),dtype=torch.int).cuda()
        self.surfacey=torch.zeros((n + [3,]),dtype=torch.int).cuda()
        self.surfacez=torch.zeros((n + [3,]),dtype=torch.int).cuda()
        self.mrollpad=torch.zeros(list(np.array(n)+2)+[3,],dtype=torch.float32).cuda()
        self.f_h_demag_pad = torch.zeros([2*i-1 for i in n] + [3,], dtype=torch.complex64).cuda()
    def hzee_(self,h_zee):
        hzee= torch.zeros(self.n + [3,],dtype=torch.float32).cuda()
        hzee[:,:,:,0]=h_zee[0]
        hzee[:,:,:,1]=h_zee[1]
        hzee[:,:,:,2]=h_zee[2]
        return hzee
    def h_dem(self,m,):
        n=self.n
        f_h_demag_pad=self.f_h_demag_pad
        self.m_pad[:n[0],:n[1],:n[2],:] = m#pad 指补零扩大矩阵
        f_m_pad =torch.fft.fftn(self.m_pad,dim=(0,1,2))
        f_h_demag_pad[:,:,:,0] = torch.sum(self.f_n_demag[:,:,:,(0, 1, 2)]*(f_m_pad),dim=3)#  N元素如👉  0  1  2 张量内积
        f_h_demag_pad[:,:,:,1] = torch.sum(self.f_n_demag[:,:,:,(1,3,4)]*(f_m_pad),dim=3)#          👉   1  3  4
        f_h_demag_pad[:,:,:,2] = torch.sum(self.f_n_demag[:,:,:,(2,4,5)]*(f_m_pad),dim=3)#          👉   2  4  5
        h_demag = torch.fft.ifftn(f_h_demag_pad,dim=(0,1,2))[:n[0],:n[1],:n[2],:].real#取实部
        return -self.ms*h_demag
    def h_eff(self,m,hzee):
        hdem=self.h_dem(m=m,)
        mrollpad=self.mrollpad
        mrollpad[1:-1,1:-1,1:-1]=m
        m1r=torch.roll(mrollpad,1,dims=0)
        m2r=torch.roll(mrollpad,1,dims=1)
        m3r=torch.roll(mrollpad,1,dims=2)
        m4r=torch.roll(mrollpad,-1,dims=0)
        m5r=torch.roll(mrollpad,-1,dims=1)
        m6r=torch.roll(mrollpad,-1,dims=2)
        hex=2*self.A/(self.mu0*self.ms)*(
        - 2 * m * sum([1/x**2 for x in self.dx])   \
        +((m1r+m4r)[1:-1,1:-1,1:-1]*self.m_form+m*self.surfacex)/(self.dx[0]**2)      \
        +((m2r+m5r)[1:-1,1:-1,1:-1]*self.m_form+m*self.surfacey)/(self.dx[1]**2)      \
        +((m3r+m6r)[1:-1,1:-1,1:-1]*self.m_form+m*self.surfacez)/(self.dx[2]**2) )
        return hex+hzee+hdem
    def rk45(self,y0,t_span, TOL, mindt, maxdt,N,h_zee,fixdt,):
        flag = True
        ti = t_span[0]
        if fixdt!=0 : dt=fixdt
        else        : dt=maxdt
        y=self.normaliza(y0)
        t_y=[]
        R_=[]
        k1=self.llg_func(ti, y)
        t_y.append([ti]+list(map(lambda i:torch.sum( y[:,:,:,i])/N, range(3)) ))
        R_.append(torch.tensor([0],dtype=torch.float32))
        step=0
        yi=torch.zeros(y.shape,dtype=torch.float32).cuda()
        while flag:
            if fixdt!=0 : dt=fixdt
            yi.copy_(y)
            h = dt*self.gamma

            k1h=h*k1
            y2=yi+0.2*k1h
            k2h=self.llg_func(ti+.2*dt, self.normaliza(y2),)*h
            
            y3=yi+3.*k1h/40.+9.*k2h/40.
            k3h=self.llg_func(ti+.3*dt, self.normaliza(y3),)*h

            y4=yi+44.*k1h/45.-56.*k2h/15.+32.*k3h/9.
            k4h=self.llg_func(ti+.4*dt,self.normaliza (y4),)*h

            y5=yi+19372.*k1h/6561.-25360.*k2h/2187.+64448.*k3h/6561.-212.*k4h/729.
            k5h=self.llg_func(ti+8./9.*dt,self.normaliza (y5),)*h

            y6=yi+9017.*k1h/3168.-355.*k2h/33.+46732.*k3h/5247.+49.*k4h/176.-5103.*k5h/18656.
            k6h=self.llg_func(ti+dt,self.normaliza (y6),)*h

            yi=self.normaliza(yi+35.*k1h/384.+500.*k3h/1113.+125.*k4h/192.-2187.*k5h/6784.+11.*k6h/84.)
            k7=self.llg_func(ti+dt,yi,)

            D = ((35./384.)-(5179./57600.))*k1h+((500./1113.)-(7571./16695.))*k3h+((125./192.)-(393./640.))*k4h+((-2187./6784.)-(-92097./339200.))*k5h+((11./84.)-(187./2100.))*k6h+((0.)-(1./40.))*k7*h
            R = torch.sqrt(torch.sum(D*D,dim=3).max())
             
            if R<=TOL or dt<=mindt or fixdt!=0: # accurate is acceptable
                ti += dt
                t_y.append([ti]+list(map(lambda i:torch.sum( yi[:,:,:,i])/N, range(3)) ))
                y=yi
                R_.append(R)
                q_arg=1/5
                k1=k7
                print(step,"  ",ti,"  ",dt)
                step+=1
            else:
                q_arg=1/6
            # compute q
            if fixdt==0:
                q=(TOL/R)**(q_arg)
                if q<=.5:
                    dt*=.5
                if q>=2:
                    dt*=2
                if q==1.:#跳出死循环
                    dt*=0.99
                else:
                    dt *= q.cpu().numpy()
                if dt>=maxdt:
                    dt=maxdt
                if dt<=mindt:
                    dt=mindt
            if ti >= t_span[1]:
                flag=False
            elif ti+dt > t_span[1]:
                dt = t_span[1]-ti
                if dt<mindt:
                    flag=False
        return [y,t_y,R_]
    def llg_func(self,t,y0):
        h_eff=self.h_eff(y0,self.hzee_(self.h_zee(t)),)
        f= - 1/(1+self.alpha**2) * torch.cross(y0, h_eff) - self.alpha/(1+self.alpha**2) * torch.cross(y0, torch.cross(y0, h_eff))
        return f
    def normaliza(self,m):
        #return m
        return m/(torch.repeat_interleave(torch.sqrt(torch.sum(m*m,dim=3 )),3).reshape(m.shape)+1e-18)
    def run(self):
        name=self.name
        n=self.n
        m0=self.m0
        h_zee=self.h_zee
        T=self.T
        T0=self.T0
        error=self.error
        mindt=self.mindt
        maxdt=self.maxdt
        fixdt=self.fixdt
        inscribed=self.inscribed
    #==========初始化===============================================
        if isinstance(m0,list):
            self.m[:,:,:,0]   = m0[0]
            self.m[:,:,:,1]   = m0[1]
            self.m[:,:,:,2]   = m0[2]
        else:
            self.m=torch.tensor(m0).cuda()
        m_form_numpy= np.ones(n + [3,],dtype=np.int)#构建形状模板
        N=np.prod(list(n))#格点个数
        if inscribed:
            N=0
            m_form_numpy=np.zeros_like(m_form_numpy)
            for i in range(int(n[0])):
                for j in range(int(n[1])):
                    for k in range(int(n[2])):
                        if ((i-((n[0]-1)/2))**2)/(((n[0]-1)/2)**2)+((j-((n[1]-1)/2))**2)/(((n[1]-1)/2)**2)+((k-((n[2]-1)/2))**2)/(((n[2]-1)/2)**2)<=1:
                            m_form_numpy[i,j,k]=1
                            N+=1
        n_demag=np.load(self.namedemag+".npy")
        f_n_demag =np.fft.fftn(n_demag,axes=[0,1,2])
        self.f_n_demag=torch.tensor(f_n_demag).cuda()
        #=====ROLL 模板 交换场======================================5=============
        self.m_form=torch.tensor(m_form_numpy,dtype=torch.int).cuda()
        self.m=self.m*self.m_form

        form1_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form1_roll[0,:,:,:]=0
        form1_roll*=self.m_form
        
        form2_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form2_roll[:,0,:,:]=0
        form2_roll*=self.m_form

        form3_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form3_roll[:,:,0,:]=0
        form3_roll*=self.m_form

        form4_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form4_roll[-1,:,:,:]=0
        form4_roll*=self.m_form

        form5_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form5_roll[:,-1,:,:]=0
        form5_roll*=self.m_form

        form6_roll=torch.ones_like(self.m_form,dtype=torch.int)
        form6_roll[:,:,-1,:]=0
        form6_roll*=self.m_form
        #======ROLL==========================================
        m1r=torch.roll(self.m_form,1,dims=0)
        m1r*=form1_roll
        m2r=torch.roll(self.m_form,1,dims=1)
        m2r*=form2_roll
        m3r=torch.roll(self.m_form,1,dims=2)
        m3r*=form3_roll
        m4r=torch.roll(self.m_form,-1,dims=0)
        m4r*=form4_roll
        m5r=torch.roll(self.m_form,-1,dims=1)
        m5r*=form5_roll
        m6r=torch.roll(self.m_form,-1,dims=2)
        m6r*=form6_roll
        #=======surface========================================
        surface1=self.m_form-m1r
        surface2=self.m_form-m2r
        surface3=self.m_form-m3r
        surface4=self.m_form-m4r
        surface5=self.m_form-m5r
        surface6=self.m_form-m6r
        surface=surface1+surface2+surface3+surface4+surface5+surface6
        self.surfacex=surface1+surface4
        self.surfacey=surface2+surface5
        self.surfacez=surface3+surface6
        #===============================================================================
        #=============LLG演化=========================
        #================================================
        print("runing")
        m,m_run,r_run=self.rk45(self.m,[T0,T],error,mindt,maxdt,N,h_zee,fixdt)
        for i in range(len(m_run)):
            m_run[i][1]=m_run[i][1].cpu().numpy()
            m_run[i][2]=m_run[i][2].cpu().numpy()
            m_run[i][3]=m_run[i][3].cpu().numpy()
            r_run[i]=r_run[i].cpu().numpy()
        np.save(str(name)+"SIM"+".npy",np.array(m_run))
        np.save(str(name)+"SIM"+"error.npy",np.array(r_run))
        np.save(str(name)+"SIMnow"+".npy",m.cpu().numpy())
#===========================================================================================
#==================================================================================================
#===============-------main-------===========================================================================================
#=======================================================================================================
if __name__ == "__main__":
    st=time.perf_counter()
    sim(    name="sp4",#所有文件都用此命名
    namedemag="sp4_ndemag",
    n=[100,25,1],#cell数目
    dx=[5e-9,5e-9,3e-9],#cell大小
    m0=[1,0.1,0],#初始磁矩，vector用list，具体分布用numpy ndarray
    gamma = 2.211e5,#旋磁比
    ms    = 800e3,#饱和磁化强度
    A     = 13e-12,#交换常数
    alpha=1,#damping
    h_zee=lambda t:[0,0,0],#外场
    T=1e-9,#总仿真时间
    T0=0,
    fixdt=1e-12,#RK45定步长      当fixdt=0时，
    #                                 maxdt，mindt，error才有意义
    #                            当fixdt！=0时，
    #                                  maxdt，mindt，error不起作用
    mindt=1e-13,#RK45最小迭代步长
    maxdt=1e-10,#RK45最大迭代步长   由于第一次迭代步长为maxdt，设置太大容易溢出
    error=1e-5,#RK45允许的误差
    inscribed=False).run()
    en=time.perf_counter()
    print(en-st)
